#!/usr/bin/python

import urllib                   # to download files
import sys                      # to write messages to buffers
import getopt                   # to handle command line options
import subprocess               # to use external unzippers
import os.path                  # to manipulate filenames
import shutil                   # file copying

from tcclimFunctions import *   # get common functions 


## FULL DATA SCRIPT HAMPERED BY THE LACK OF LZW FILE DECOMPRESSION FOR PYTHON
## BUT 7ZIP provides a windows solution for .Z and pax archives - maybe also Mac/UNIX, 
## but pax/uncompress also available for those platforms

## ~ 15 minutes including download time on a standard Windows 7 desktop

def usage():
  
  use = """
  
  tcclimDownloadCRUTS2.py:
  
  Downloads component data files for the CRU TS 2.10 climate
  datasets and creates transposed data files for use by the
  tcclimCruTS2.py program. The transposed files will be saved
  in the folder 'tcclim_cru_ts_2_10'.
  
  Basic use:
     tcclimDownloadCRUTS2.py
  
  Options:
     -h or --help: show this help information
     -k or --keep: retain the original untransposed data files
     -7 or --7zip: path to the 7zip program on Windows
     -u or --uncompress: path to the uncompress program on Mac or Linux
  
  This program requires the use of external programs to extract
  the original data from compressed file formats. On Windows, 
  the open-source program 7zip is used and must be installed 
  before using this program. By default, the 7zip program is 
  expected to be found at 'C:\Program Files\7Zip\7Z.exe', but 
  this may be changed using the --7zip option. On Mac or Linux
  uncompress is used and the path to this may also be specified.
  
  """
  
  sys.stdout.write(use)


def main():
    
    # sort out the options and argument
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hks:u:", ["help","keep","sevenzip","uncompress"])
    except getopt.GetoptError, err:
        # print help information and exit:
        sys.stderr.write("\n  " + str(err) + '\n') # will print something like "option -a not recognized"
        usage()
        sys.exit()
    
    # initial setups
    keep = False
    sevenzip = r"C:\Program Files\7-Zip\7Z.exe"
    uncompress = 'uncompress'
    
    for o, a in opts:
        if o in ('-h', '--help'):
            usage()
            sys.exit()
        elif o in ('k', '--keep'):
            keep = True
        elif o in ('s', '--sevenzip'):
            sevenzip = a
        elif o in ('u', '--uncompress'):
            uncompress = a
        else:
            assert False, "Unhandled option"
            usage()
            sys.exit()
    
    # Pre flight checks:
    #  - for utilities 
    if os.name == 'posix':
        # use which to test for uncompress
        try:  
           proc = subprocess.Popen('which ' + uncompress ,
                                   shell=True, stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE)
           procComm = proc.communicate() # get the stdout and stderr
           # the stdout of procComm will be an empty string if the executable is not found
           if procComm[0] == '':
               sys.stderr.write('\nuncompress not found: ' + uncompress + '\n\n')
               usage()
               sys.exit() 
        except OSError, e:
           sys.stderr.write("Execution failed:" + e)
    elif os.name == 'nt':
        if not os.path.exists(sevenzip):
            sys.stderr.write('\n7Zip program not found\n\n')
            usage()
            sys.exit()
    
    # - for data directories
    ## create data directories
    if not os.path.exists('cru_ts_2_10'):
        os.mkdir('cru_ts_2_10')
    else:
        sys.stderr.write('Download folder \'cru_ts_2_10\' already exists.\n\n')
        sys.exit()
    
    if not os.path.exists('tcclim_cru_ts_2_10'):
        os.makedirs('tcclim_cru_ts_2_10')
    else:
        sys.stderr.write('\nTransposed file directory \'tcclim_cru_ts_2_10\' already exists...\n\n')
        sys.exit()
    
    ##########################
    ## DOWNLOAD CRU TS 2.10 ##
    ##########################
    
    cruts2repos = 'http://www.cru.uea.ac.uk/cru/data/hrg/cru_ts_2.10/data_all/'
    
    cruts2files = ['cru_ts_2_10.1901-2002.cld.Z' ,'cru_ts_2_10.1901-2002.dtr.Z',
                   'cru_ts_2_10.1901-2002.frs.Z','cru_ts_2_10.1901-2002.pre.Z',
                   'cru_ts_2_10.1901-2002.tmn.Z','cru_ts_2_10.1901-2002.tmp.Z',
                   'cru_ts_2_10.1901-2002.tmx.Z','cru_ts_2_10.1901-2002.vap.Z',
                   'cru_ts_2_10.1901-2002.wet.Z']
    
    sys.stdout.write('\n\nDownloading CRU TS 2.10     [.....................]\n')
    
    for eachF in cruts2files:
        remotename = cruts2repos + eachF
        localname  = os.path.join('cru_ts_2_10', eachF)
        sys.stdout.write(eachF + ' [')
        sys.stdout.flush()
        try:
            urllib.urlretrieve(remotename, localname, progressTicker)
        except IOError, e[1][1]:
            sys.stderr.write("\n\nFile " + eachF + " could not be downloaded:\n" + e + "\nCheck network connection.\n\n")
            sys.exit()
        sys.stdout.write(']\n')
    
    sys.stdout.write('\nCRU TS 2.1 dataset download complete.\n\n')
    
    ##########################
    ## EXTRACT CRU TS  2.10 ##
    ##########################
    
    os.chdir('cru_ts_2_10')
    
    # do one thing on windows and another on posix
    # use os.name to get 'nt' (Windows) versus 'posix' (Darwin/Linux)
    # - could use sys.platform but more variation in mac/linux returns
    if os.name == 'posix':
        # use uncompress
        for eachF in cruts2files:
            sys.stdout.write('Uncompressing ' + eachF + '\n')
            sys.stdout.flush()
            try:  
               proc = subprocess.Popen(uncompress + ' ' + eachF,
                                       shell=True, stderr=subprocess.PIPE,
                                       stdout=subprocess.PIPE)
               procWait = proc.wait()        # hang on until the extraction ends
               procComm = proc.communicate() # get the stdout and stderr
            except OSError, e:
               sys.stderr.write("Execution failed:" + e)
    elif os.name == 'nt':
       # use 7zip
       for eachF in cruts2files:
           sys.stdout.write('Uncompressing ' + eachF + '\n')
           sys.stdout.flush()
           try:  
              proc = subprocess.Popen('"' + sevenzip + '" e ' + eachF,
                                      shell=True, stderr=subprocess.PIPE,
                                      stdout=subprocess.PIPE)
              procWait = proc.wait()        # hang on until the extraction ends
              procComm = proc.communicate() # get the stdout and stderr
           except OSError, e:
              sys.stderr.write("Execution failed:" + e)
    
    ###########################
    ## TRANSPOSE CRU TS 2.10 ##
    ###########################
    
    ## TRANSPOSE AN EXISTING COPY OF THE CRU TS 2 DATASET
    ## - check directories
    
    os.chdir('../')
    sys.stdout.write('\nTransposing data to \'tcclim_cru_ts_2_10\'.\n\n')
    
    ################################################
    ## translate variables: observed 1901-2000    ##
    ## (9 variables)                              ##
    ################################################
    
    for eachF in cruts2files:
        # transpose
        eachF = eachF[:-2] # remove .Z extension
        infile  = 'cru_ts_2_10/' + eachF
        outfile = 'tcclim_cru_ts_2_10/tr.' + eachF
        sys.stdout.write( 'Transposing ' + eachF + '\n')
        try:
            transposeGrim(infile, outfile)
        except:
            sys.stderr.write( 'Skipping missing file: ' + eachF + '\n')
    
    ###################################################
    ## create a file containing grid-refs in the     ##
    ## order in which they appear in the grim files  ##
    ###################################################
    
    # scan last input grim file for grid refs
    getGridRefIndexGrim(infile, 'tcclim_cru_ts_2_10/transposedGridRefIndex.txt')
    
    if keep:
        sys.stdout.write( '\n\nCRU TS 2.10 transposition completed.\nOriginal data directory retained.\n\n')
    else:
        shutil.rmtree('cru_ts_2_10/')
        sys.stdout.write( '\n\nCRU TS 2.10 transposition completed.\nOriginal data directory deleted.\n\n')


if __name__ == "__main__":
    main()

